import argparse
from sys import platform
from models2 import *
from utils.datasets import *
from utils.utils import *
import pickle

import numpy as np
# np.set_printoptions(threshold=10000)
# torch.set_printoptions(profile="full")
def detect(save_txt=False, save_img=False):
    img_size = opt.img_size  # (320, 192) or (416, 256) or (608, 352) for (height, width)
    out, source, weights= opt.output, opt.source, opt.weights

    # Initialize
    device = torch_utils.select_device(opt.device)
    if os.path.exists(out):
        shutil.rmtree(out)  # delete output folder
    os.makedirs(out)  # make new output folder

    # Initialize model
    model = Darknet(opt.cfg, img_size, quan=opt.quan, save_act=opt.save_act, conv_bias=not opt.no_bias)

    # Load weights
    if weights.endswith('.pt'):  # pytorch format
        model.load_state_dict(torch.load(weights, map_location=device)['model'])
    else:
        load_darknet_weights(model, weights)
    
    # Eval mode
    model.to(device).eval()
    dataset = LoadImages(source, img_size=img_size, half=False)

    # Get names and colors
    names = load_classes(opt.names)
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
    
    if opt.save_wgt:
        wgt_parm={}
        for name,parameters in model.named_parameters():
            if 'Conv2d' in name:
                wgt_parm[name]=parameters.cpu().detach().numpy()        
        import pickle
        f_wgt=open(opt.out_wgt,'wb')
        pickle.dump(wgt_parm, f_wgt)
        f_wgt.close()
    
    if opt.prune:
        prune_weight(model, opt.threshold)
    
    if opt.quan:
        quan_weight(model)
    
    # Run inference
    for path, img, im0s, vid_cap in dataset:
        # cv2.imwrite("tmp.jpg", img.transpose(1, 2, 0)*255)
        # Get detections
        img = torch.from_numpy(img).to(device)
        
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        print(img)
        pred = model(img)[0]
        # Apply NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.nms_thres)
        print(len(pred))
        print('hhh')
        # Process detections
        for i, det in enumerate(pred):  # detections per image
            p, s, im0 = path, '', im0s

            save_path = str(Path(out) / Path(p).name)
            s += '%gx%g ' % img.shape[2:]  # print string
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += '%g %ss, ' % (n, names[int(c)])  # add to string

                # Write results
                for *xyxy, conf, cls in det:
                    label = '%s %.2f' % (names[int(cls)], conf)
                    plot_one_box(xyxy, im0, label=label, color=colors[int(cls)])

            # Save results (image with detections)
            cv2.imwrite(save_path, im0)
        
        if opt.save_act:
            model.save_act_pkl(opt.out_act)

    

    print('Results saved to %s' % os.getcwd() + os.sep + out)
    # print('Done. (%.3fs)' % (time.time() - t0))
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='cfg/yolov3_ReLU-voc.cfg', help='*.cfg path')
    parser.add_argument('--names', type=str, default='data/voc.names', help='*.names path')
    parser.add_argument('--weights', type=str, default='/home/zhangjm/backup/yolov3_ReLU_voc/yolov3_ReLU-voc_final.weights', help='path to weights file')
    parser.add_argument('--source', type=str, default='/home/zhangjm/yolov3_voc/data/samples/dog.jpg', help='source')  # input file/folder, 0 for webcam
    parser.add_argument('--output', type=str, default='output', help='output folder')  # output folder
    parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
    parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression')

    parser.add_argument('--device', default='', help='device id (i.e. 0 or 0,1) or cpu')
    parser.add_argument('--quan', action='store_true', help='quantize the model')
    parser.add_argument('--no_bias', action='store_true', help='whether model has bias')
    parser.add_argument('--save-act', action='store_true', help='save activation')
    parser.add_argument('--save-wgt', action='store_true', help='save wgt')
    parser.add_argument('--out-act', type=str, default='/home/zhangjm/backup/20210321/act.pkl', help='save activation')
    parser.add_argument('--out-wgt', type=str, default='/home/zhangjm/backup/20210321/wgt.pkl', help='save wgt')
    parser.add_argument('--prune', action='store_true', help='prune the model')
    parser.add_argument('--threshold', type=float, default=0.005)
    
    opt = parser.parse_args()
    print(opt)

    with torch.no_grad():
        detect()
